# You need to set the environmental variable OPENAI_API_KEY
from langchain_core.tools import tool
import streamlit as st
from langchain_zhipu import ChatZhipuAI

from langchain.agents import create_react_agent
from langchain.agents import AgentType
from langchain.agents import AgentExecutor

from langchain_core.prompts import PromptTemplate
from langchain.agents.output_parsers import ReActJsonSingleInputOutputParser
from langchain.tools.render import render_text_description_and_args
from customer_logging import get_logger
from data.user_data import obtain_data
from ChatbotWithRetrieval import ChatbotWithRetrieval
from datetime import datetime

#
# 问题：XX订单的回寄地址
# 期望ai能回复：这个订单对应售后单的回寄地址信息。
# 问题：XX订单的换货单号有了吗
# 期望ai能回复：这个订单对应售后单的寄出快递单号。
logger = get_logger("sys-agent=")


def get_exchange_note_info(exchange_note_number_list):
    exchange_note_info_list = []
    for return_address in exchange_note_number_list:
        order_after_no = return_address['order_after_no']
        tracking_to_company = return_address['tracking_to_company']
        tracking_to_num = return_address['tracking_to_num']
        exchange_note_info = f"售后单：{order_after_no}   商户寄出物流公司(换货单公司)：{tracking_to_company}  商户寄出物流单号(换货单号)：{tracking_to_num}"
        exchange_note_info_list.append(exchange_note_info)
    return exchange_note_info_list


def get_return_address(return_address_list):
    return_address_info_list = []
    for return_address in return_address_list:
        order_after_no = return_address['order_after_no']
        province_name = return_address['province_name']
        city_name = return_address['city_name']
        district_name = return_address['district_name']
        address = return_address['address']
        name = return_address['name']
        mobile = return_address['mobile']
        tracking_company = get_with_default(return_address, 'tracking_company')
        tracking_num = get_with_default(return_address, 'tracking_num')
        return_address_info = f"售后单：{order_after_no}   {province_name}/{city_name}/{district_name}/{address} 姓名：{name}  手机：{mobile}   物流公司：{tracking_company}  寄回物流单号：{tracking_num}"
        return_address_info_list.append(return_address_info)
    return return_address_info_list


# Custom tool for the Agent

def get_with_default(d, key, default="无"):
    return d.get(key) or default


@tool
def obtain_return_address_for_order(order_no):
    """
    This method obtains the return address of the order based on the order number,Not considering historical information, receive value such as 20200831163949000053,
    return "售后单：2020090216443300007   浙江省/杭州市/滨江区/吉利大厦 姓名：杨小雨  手机：17306409987   物流公司：顺丰  寄回物流单号：3335"
    """
    return_address = obtain_data(f"""
    SELECT
        order_after_no,
        province_name ,
        city_name ,
        district_name ,
        address,
        name,
        mobile,
        tracking_company,
        tracking_num
    FROM
        order_after_sale 
    WHERE
        order_no = {order_no} having province_name != '' """)
    logger.info("obtain_return_address_for_order=order_no:" + order_no)
    if return_address is None:
        return "暂无退货/回寄地址，请检查对应售后单"
    return get_return_address(return_address)
    # return "杭州市滨江区长河街道165号 李先生  手机：15432314563"


# Custom tool for the Agent
@tool
def obtain_the_exchange_note_info_for_order(order_no):
    """
    This method is to obtain the exchange note info for the order,Not considering historical information, receive value such as 20200831163949000053,
    return 售后单：2020090216443300007   商户寄出物流公司(换货单公司)：韵达  商户寄出物流单号(换货单号)：JT3066225372944 or 73521050381737  or YDccc
    """

    exchange_note_number_list = obtain_data(f"""
    SELECT
        order_after_no,
        tracking_to_num,
        tracking_to_company 
    FROM
        order_after_sale 
    WHERE
        order_no = {order_no} having tracking_to_num != '' """)
    logger.info("obtain_return_address_for_order=order_no:" + order_no)
    if exchange_note_number_list is None:
        return f"暂未查询到订单:{order_no} 换货单号"
    return get_exchange_note_info(exchange_note_number_list)


@tool
def obtain_product_or_general_question_answer(question):
    """
This method is used to query product or general question and answer information.
If Action=None or other methods are not applicable, use the result set of this method as the final answer.
"""
    logger.info("obtain_knowledge_base_information:" + str(question))
    if type(question) == dict:
        question = question['title']

    logger.info("obtain_knowledge_base_information:" + question)
    llm = get_llm()
    chat_chain = llm.qa
    # 获取当前时间
    now = datetime.now()
    # 格式化当前时间为 yyyy-MM-DD 格式
    formatted_now = now.strftime('%Y-%m-%d')
    return chat_chain.invoke({"history": None, "input": question, "current_time": formatted_now})["answer"]

    # return "商城要求供应商这边是48小时内安排发货的（快递停发地区等特殊情况除外），具体时间以系统发货时间为准哦，请耐心等待! "


@st.cache_resource
def get_llm():
    logger.info("加载对话大模型")
    base_dir = 'One'
    llm = ChatbotWithRetrieval(base_dir)
    return llm


@st.cache_resource
def obtain_answer():
    logger.info("枫华 agent")
    agent = FhAgent()
    agent_executor = agent.agent_executor
    return agent_executor


class FhAgent:

    def __init__(self):
        template = '''
        SYSTEM

        Answer the following questions as best you can. You have access to the following tools:



        {tools}



        The way you use the tools is by specifying a json blob.

        Specifically, this json should have a `action` key (with the name of the tool to use) and a `action_input` key (with the input to the tool going here).



        The only values that should be in the "action" field are: {tool_names}



        The $JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. Here is an example of a valid $JSON_BLOB:



        ```

        {{

          "action": $TOOL_NAME,

          "action_input": $INPUT

        }}

        ```



        ALWAYS use the following format:



        Question: the input question you must answer

        Thought: you should always think about what to do

        Action:

        ```

        $JSON_BLOB

        ```

        Observation: the result of the action

        ... (this Thought/Action/Observation can repeat N times)

        Thought: I now know the final answer

        Final Answer: the final answer to the original input question



        Begin! Reminder to always use the exact characters `Final Answer` when responding.
        The tool may return multiple lines of information, please return all lines when returning.
        If the tool returns the final answer, return it directly in the original information and its format, and do not omit the information returned by the tool.
        Previous conversation history:

        {chat_history}

        New HUMAN input: {input}



        {agent_scratchpad}'''

        prompt = PromptTemplate.from_template(template)
        model = ChatZhipuAI(api_key="a6d0c5ecfb1376673f0500b71fd46443.eDfhQqZkP6cnAkcw"
                            , model="glm-4", verbose=True)
        model.do_sample = False  # 温度设置为0，结果随机性 ghbnm

        tools = [obtain_return_address_for_order, obtain_the_exchange_note_info_for_order,
                 obtain_product_or_general_question_answer]
        agent = create_react_agent(
            model,
            tools,
            prompt,
            tools_renderer=render_text_description_and_args,
            output_parser=ReActJsonSingleInputOutputParser(),
        )
        # ,max_iterations=2
        self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
    # answer = agent_executor.invoke({"input": "订单20200902152124000028的回寄地址是多少？", "history": ["user:订单32123的换货单号有了么？"]})
    # answer = agent_executor.invoke({"input": "订单2545的回寄地址是多少？"})
    # answer = agent_executor.invoke({"input": "订单20200902152124000028的换货单号有了么？"})
    # answer = agent_executor.invoke({"input": "订单32123的换货单号有了么？"})
    # answer = agent_executor.invoke({"input": "什么时间发货"})
    # print(answer)
